# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

from recordclass import recordclass

from coin.proto.coin_order_gateway_pb2 import (OrderEvent,
                                               OrderSubmitRequestProto,
                                               OrderAmendRequestProto,
                                               OrderCancelRequestProto)
import coin.proto.coin_order_enums_pb2 as coin_order
import coin.proto.coin_market_enums_pb2 as coin_enum

OrderEventStats = recordclass('OrderEventStats',
                              [
                                  'internal_order_id',
                                  'accept_count',
                                  'accept_qty',
                                  'fill_count',
                                  'fill_qty',
                                  'is_partial_fill',
                                  'submit_reject_count',
                                  'submit_reject_qty',
                                  'submit_qty',
                                  'submit_ts',
                                  'accept_ts',
                                  'fill_ts',
                                  'cancel_ts',
                                  'cancel_confirmed_ts'
                              ])

Submit2AcceptStats = recordclass('Submit2AcceptStats',
                                 ['accept_count', 'submit2accept_delay_total'])

Submit2FillStats = recordclass(
    'Submit2FillStats',
    ['immediate_fill_count', 'accept_then_fill_count', 'submit2fill_delay_total'])


def create_empty_stats(internal_order_id):
  return OrderEventStats(internal_order_id,
                         0,
                         0,
                         0,
                         0,
                         False,
                         0,
                         0,
                         0,
                         None,
                         None,
                         None,
                         None,
                         None)


class OrderEventStatsCollector(object):
  def __init__(self):
    self.order_stats = {}

  def update_by_submit_request(self, submit_request: OrderSubmitRequestProto, timestamp):
    internal_order_id = submit_request.internal_order_id
    order_stats = self.order_stats.get(internal_order_id, [])
    order_stats_item = create_empty_stats(internal_order_id)
    order_stats_item.submit_ts = timestamp
    order_stats.append(order_stats_item)
    self.order_stats[internal_order_id] = order_stats

  def update_by_amend_request(self, amend_request: OrderAmendRequestProto, timestamp):
    internal_order_id = amend_request.internal_order_id
    order_stats = self.order_stats.get(internal_order_id, None)
    assert order_stats is not None
    order_stats_item = create_empty_stats(internal_order_id)
    order_stats_item.submit_ts = timestamp
    order_stats.append(order_stats_item)

  def update_by_cancel_request(self, cancel_request: OrderCancelRequestProto, timestamp):
    if not cancel_request.HasField('internal_order_id'):
      return
    internal_order_id = cancel_request.internal_order_id
    order_stats = self.order_stats.get(internal_order_id, None)
    assert order_stats is not None, internal_order_id
    order_stats[-1].cancel_ts = timestamp

  def update_by_order_event(self, order_event: OrderEvent):
    if not order_event.HasField('internal_order_id'):
      return
    internal_order_id = order_event.internal_order_id
    order_stats = self.order_stats.get(internal_order_id, None)
    assert order_stats is not None, internal_order_id
    order_stats_item = order_stats[-1]
    if order_event.type in [OrderEvent.ORDER_SUBMITTED, OrderEvent.ORDER_AMENDED]:
      order_stats_item.submit_qty = order_event.order_qty
    elif order_event.type in [OrderEvent.ORDER_ACCEPTED, OrderEvent.ORDER_AMEND_ACCEPTED]:
      # Duplicate accept message could happen.
      if order_stats_item.accept_count == 0:
        order_stats_item.accept_count += 1
        order_stats_item.accept_qty += order_event.order_qty
        order_stats_item.accept_ts = order_event.event_time
        order_stats_item.is_partial_fill = False
    elif order_event.type == OrderEvent.ORDER_FILLED:
      order_stats_item.fill_qty += order_event.fill_qty
      if order_event.exchange_type == coin_enum.Bitflyer:
        order_stats_item.fill_count += 1
        order_stats_item.fill_ts = order_event.event_time
      elif order_event.exchange_type in [coin_enum.Okex, coin_enum.Bitmex]:
        if not order_stats_item.is_partial_fill:
          order_stats_item.fill_count += 1
          order_stats_item.fill_ts = order_event.event_time
        if order_event.order_state == coin_order.WORKING_ORDER:
          order_stats_item.is_partial_fill = True
      else:
        raise ValueError(order_event.exchange_type)
    elif order_event.type == OrderEvent.CANCEL_CONFIRMED:
      order_stats_item.cancel_confirmed_ts = order_event.event_time
    elif order_event.type == OrderEvent.ORDER_REJECTED:
      if order_event.order_event_source == OrderEvent.EXCHANGE:
        order_stats_item.submit_reject_count += 1
        order_stats_item.submit_reject_qty += order_stats_item.submit_qty
    else:
      pass

  def calc_delay_stats(self):
    submit2accept_stats = Submit2AcceptStats(0, 0)
    submit2fill_stats = Submit2FillStats(0, 0, 0)
    for stats in self.order_stats.values():
      for item in stats:
        if item.submit_ts and item.fill_ts:
          if not item.accept_ts:
            # Immediate fill case.
            submit2accept_stats.accept_count += 1
            submit2accept_stats.submit2accept_delay_total += \
                item.fill_ts - item.submit_ts
            submit2fill_stats.immediate_fill_count += 1
            submit2fill_stats.submit2fill_delay_total += item.fill_ts - item.submit_ts
          else:
            # Accept then fill case.
            submit2fill_stats.accept_then_fill_count += 1
            submit2fill_stats.submit2fill_delay_total += item.fill_ts - item.submit_ts
        if item.submit_ts and item.accept_ts:
          submit2accept_stats.accept_count += 1
          submit2accept_stats.submit2accept_delay_total += \
              item.accept_ts - item.submit_ts
    return submit2accept_stats, submit2fill_stats

  @property
  def total_accept_count(self):
    total_accept_count = 0
    for stats in self.order_stats.values():
      for item in stats:
        assert item.accept_count <= 1
        if item.accept_count == 0 and item.fill_count == 1:
          total_accept_count += 1
        total_accept_count += item.accept_count
    return total_accept_count

  @property
  def total_accept_qty(self):
    total_accept_qty = 0
    for stats in self.order_stats.values():
      for item in stats:
        if item.accept_count == 0 and item.fill_count == 1:
          total_accept_qty += item.fill_qty
        total_accept_qty += item.accept_qty
    return total_accept_qty

  @property
  def total_fill_count(self):
    total_fill_count = 0
    for stats in self.order_stats.values():
      for item in stats:
        total_fill_count += item.fill_count
    return total_fill_count

  @property
  def total_fill_qty(self):
    total_fill_qty = 0
    for stats in self.order_stats.values():
      for item in stats:
        total_fill_qty += item.fill_qty
    return total_fill_qty

  @property
  def total_submit_count(self):
    total_submit_count = 0
    for stats in self.order_stats.values():
      for _ in stats:
        total_submit_count += 1
    return total_submit_count

  @property
  def total_submit_qty(self):
    total_submit_qty = 0
    for stats in self.order_stats.values():
      for item in stats:
        total_submit_qty += item.submit_qty
    return total_submit_qty

  @property
  def total_submit_reject_count(self):
    total_submit_reject_count = 0
    for stats in self.order_stats.values():
      for item in stats:
        total_submit_reject_count += item.submit_reject_count
    return total_submit_reject_count

  @property
  def total_submit_reject_qty(self):
    total_submit_reject_qty = 0
    for stats in self.order_stats.values():
      for item in stats:
        total_submit_reject_qty += item.submit_reject_qty
    return total_submit_reject_qty
